package com.cat.greedyAlgorithm;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;

/**
 * @author 曲大人的喵
 * @description https://leetcode.cn/problems/sum-of-perfect-square-ancestors/
 * @create 2025/10/12 21:13
 * @since JDK17
 */

public class Solution91 {
    List<Integer>[] g;
    int[] nums;
    HashMap<Integer, Integer> map = new HashMap<>();
    static int N = ((int) (1e5 + 1));
    static int[] core = new int[N];
    long ans;
    static {
        for (int i = 1; i < N; i++) {
            if (core[i] != 0) { // i含完全平方因子，不能作为core值
                continue;
            }
            for (int j = 1; i * j * j < N; j++) {  //
                core[i * j * j] = i;
            }
        }
    }

    void dfs(int cur) {
        map.merge(core[nums[cur]], 1, Integer::sum);
        for (int next : g[cur]) {
            int num = core[nums[next]];  //
            ans += map.getOrDefault(num, 0);    //
            // System.out.println(num);
            dfs(next);
        }
        map.merge(core[nums[cur]], -1, Integer::sum);
    }
    public long sumOfAncestors(int n, int[][] edges, int[] nums) {
        g = new List[n];
        this.nums = nums;
        Arrays.setAll(g, i -> new ArrayList<>());
        for (var e : edges) {
            g[e[0]].add(e[1]);
        }
        dfs(0);
        return ans;
    }
}
